#! /usr/bin/python
# encoding=utf-8
# Created by Fenglu Niu on 2025/5/6 21:34
from datetime import datetime

from pydantic import ConfigDict
from sqlmodel import SQLModel, Field

from common.snowflake import Snowflake

name = "entity"

# 序列化
model_config = ConfigDict(
    from_attributes=True,
    json_encoders={datetime: lambda dt: '' if dt is None else dt.strftime("%Y-%m-%d %H:%M:%S"),
                   int: lambda o: '' if o is None else o if o < 0x20000000000000 else str(o)}

)


class Entity(SQLModel):
    model_config = model_config
    id: int = Field(default_factory=Snowflake(datacenter_id=1, worker_id=1).get_id, primary_key=True)
    create_time: datetime = Field(default=datetime.now())
    update_time: datetime = Field(default=datetime.now())
    status: int = Field(default=1)
    is_deleted: int = Field(default=0)


class UpdateEntity(SQLModel):
    id: int = Field(default=None)
    update_time: datetime = Field(default=datetime.now())


class DeleteEntity(SQLModel):
    id: int = Field(default=None)
    is_deleted: int = Field(default=1)


class QueryEntity(SQLModel):
    id: int = Field(default=None)
    is_deleted: int = Field(default=0)
    start_time: datetime = Field(default=None)
    end_time: datetime = Field(default=None)
    key: str = Field(default=None)
    status: int = Field(default=None)


# 会话
class Conversation(Entity, table=True):
    """
    会话实体
    """
    __tablename__ = "hk_conversation"
    title: str = Field(default=None)
    stu_no: str = Field(default=None)
    agents: str = Field(default=None)

class ConversationQuery(QueryEntity):
    """
    会话查询实体
    """
    stu_no: str = Field(default=None)

class Message(Entity, table=True):
    """
    消息实体
    """
    __tablename__ = "hk_message"
    conversation_id: int = Field(default=None)
    role: str = Field(default=None)
    content: str = Field(default=None)
    tag: str = Field(default=None)

class MessageQuery(QueryEntity):
    """
    消息查询实体
    """
    conversation_id: int = Field(default=None)
    tag: str = 'history'